
/*
 * sndpipe: A UNIX pipes-based audio stream processor toolkit.
 * Copyright (C) 2011  Bradley Worley
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to:
 *
 *   Free Software Foundation, Inc.
 *   51 Franklin Street, Fifth Floor
 *   Boston, MA  02110-1301, USA.
 */

/* include the sndpipe header. */
#include "sndpipe.h"

/* define the fourier transform sign convention. */
#define ISIGN_FWD  1.0
#define ISIGN_REV -1.0

/* fft_next_pow2: finds the next power of two from a given number.
 * @n: the given number.
 */
unsigned long fft_next_pow2 (unsigned long n) {
  /* declare required variables. */
  unsigned long i;

  /* loop until we find the next power of two. */
  for (i = 2; i < n; i *= 2);

  /* return the calculated value. */
  return i;
}

/* fft: computes a radix-2 forward fast-fourier transform of real data.
 * @p_re: pointer to the real array.
 * @p_im: pointer to the imaginary array.
 * @N: the number of sampled data points.
 *
 * this function computes the fast fourier transform on an input array
 * of real data, sampled at discrete and evenly spaced time steps. the
 * number of input data points must be a power of two, or the function
 * will fail to return a result.
 */
int fft (double **p_re, double **p_im, unsigned long N) {
  /* declare required variables. */
  unsigned long i, j, m, n, k, step;
  double *re, *im, theta, temp;
  double ctemp[2], w[2], resi[2], resik[2];

  /* ensure we were given power-of-two sized arrays. */
  if (N & (N - 1)) {
    /* print a warning message and exit the function. */
    warnf ("unable to compute fft of non-power of two arrays");
    return SP_FAIL;
  }

  /* gain references to the output arrays. */
  re = *p_re;
  im = *p_im;

  /* left-shift the points count. */
  n = N << 1;
  j = 0;

  /* presort the points to prepare for the cooley-tukey algorithm. */
  for (i = 0; i < N - 1; i++) {
    /* swap values according to the presort rules. */
    if (j > i) {
      /* swap both the real and imaginary values. */
      SWAP (re[j], re[i]);
      SWAP (im[j], im[i]);
    }

    /* right-shift the points count once. */
    m = N >> 1;

    /* loop until we've gone through this half of the array. */
    while (m <= j) {
      /* some trickery here... :P */
      j -= m;
      m >>= 1;
    }

    /* increment the second-half array counter. */
    j += m;
  }

  /* initialize the transform loop counter. */
  k = 1;

  /* loop through the presorted data points. */
  do {
    /* calculate the step value. */
    step = 2 * k;

    /* loop through the current segment of the array. */
    for (m = 0; m < k; m++) {
      /* calculate the twiddle angle. */
      theta = M_PI * (double) m / (double) k;

      /* calculate the twiddle factors. */
      w[RE] = cos (theta);
      w[IM] = sin (theta);

      /* loop through the other segment of the array. */
      for (i = m; i < N; i += step) {
        /* store the complex number at (i+k). */
        resik[RE] = re[i + k];
        resik[IM] = im[i + k];

        /* store the complex number at (i). */
        resi[RE] = re[i];
        resi[IM] = im[i];

        /* complex-multiply the (i+k) values with the twiddle factors. */
        ctemp[RE] = w[RE] * resik[RE] - w[IM] * resik[IM];
        ctemp[IM] = w[IM] * resik[RE] + w[RE] * resik[IM];

        /* subtract the product above from the (i) value. */
        resik[RE] = resi[RE] - ctemp[RE];
        resik[IM] = resi[IM] - ctemp[IM];

        /* add the (i) value to the product, store in the (i) value. */
        resi[RE] += ctemp[RE];
        resi[IM] += ctemp[IM];

        /* store the result at (i+k). */
        re[i + k] = resik[RE];
        im[i + k] = resik[IM];

        /* store the result at (i). */
        re[i] = resi[RE];
        im[i] = resi[IM];
      }

      /* set the new index value. */
      k = step;
    }
  } while (k < N);

  /* return a success code. */
  return SP_OK;
}

/* func_fft: initiates the FFT pipes functionality.
 *
 * this function opens up an input and output pipe, reads the specified
 * channels from the incoming data, and outputs fourier-transformed data
 * on the output pipe. channels not specified by the --channels argument
 * will neither be transformed nor written to the output pipe.
 */
void func_fft (void) {
  /* declare required variables. */
  unsigned long c, C, n, N, i, k, imax, rollover, frames, *channels;
  struct pipes *pi, *po;
  double **rbuf, **wbuf, **localbuf, dmax;
  SF_INFO sfi;

  /* get the specified channels requested for transformation. */
  channels = pipes_get_channels (&C);

  /* ensure the channels option was supplied and successfully retrieved. */
  if (!channels || !C) {
    /* print a failure message and exit the application. */
    failf ("unable to retrieve target channels for fourier transform");
  }

  /* zero out the sndfile information structure. */
  memset (&sfi, 0, sizeof (sfi));

  /* attempt to open the input pipe. */
  pi = pipes_open_input ();

  /* determine whether the input pipe was opened successfully. */
  if (!pi) {
    /* print a failure message and exit the application. */
    failf ("unable to open input pipe");
  }

  /* loop through the input channels array. */
  for (c = 0; c < C; c++) {
    /* ensure that the passed channel index is within range. */
    if (channels[c] > pi->sfi.channels) {
      /* print a failure message and exit the application. */
      failf ("target channel index '%lu' out of range", channels[c]);
    }
  }

  /* set the new format information from input information. */
  sfi.samplerate = pi->sfi.samplerate;
  sfi.channels = C + 1;
  sfi.format = pi->sfi.format;

  /* ensure we have a valid output pipes format. */
  if (!sfi.format) {
    /* print a failure message and exit the application. */
    failf ("unable to retrieve pipes output format");
  }

  /* open the output pipes structure. */
  po = pipes_open_output (sfi);

  /* allocate the pipes input and output buffer memories. */
  frames = 0;
  rbuf = pipes_buf_alloc (pi, &frames);
  wbuf = pipes_buf_alloc (po, &frames);

  /* allocate local buffer memory for fft operations. */
  N = fft_next_pow2 (frames);
  localbuf = pipes_buf_alloc (po, &N);

  /* initialize the frame counter. */
  rollover = SP_FALSE;
  n = k = 0;

  /* loop while we can still read data from the input pipe. */
  while (pipes_read (pi, &rbuf, frames) == SP_OK) {
    /* check whether the local buffer is large enough for the new frames. */
    if (n + frames > N) {
      /* resize the local buffer to the next power of two number. */
      N = fft_next_pow2 (n + frames);
      localbuf = pipes_buf_realloc (&localbuf, po, N);
    }

    /* loop through the new input frames. */
    for (i = 0; i < frames; i++) {
      /* store the timebase channel value in the local buffer. */
      localbuf[0][n + i] = 0.0;

      /* loop through the input channels. */
      for (c = 0; c < C; c++) {
        /* store the input channel value in the local buffer. */
        localbuf[c + 1][n + i] = rbuf[channels[c] - 1][i];
      }

      /* check if this frame is a timebase trigger frame. */
      if (rbuf[0][i] > 0.5 && n) {
        /* this frame is a trigger frame. we've stored a whole block. */
        rollover = SP_TRUE;
        k = i;
        break;
      }
    }

    /* see if this loop iteration has stored a whole block. */
    if (rollover) {
      /* loop through the remaining local buffer frames. */
      for (i = k + 1; i < frames; i++) {
        /* and through the channels of each frame. */
        for (c = 0; c < C + 1; c++) {
          /* zero the local buffer frame value. */
          localbuf[c][n + i] = 0.0;
        }
      }

      /* loop through the output data channels. */
      for (c = 0; c < C; c++) {
        /* fourier transform the channel. */
        fft (&localbuf[c + 1], &localbuf[0], N);

        /* loop through the local buffer frames. */
        for (imax = 0, dmax = localbuf[c + 1][imax], i = 0; i < N; i++) {
          /* store the magnitude of the bin. */
          localbuf[c + 1][i] =
            sqrt (pow (localbuf[c + 1][i], 2.0)
                + pow (localbuf[0][i], 2.0));

          /* zero the imaginary bin. */
          localbuf[0][i] = 0.0;

          /* see if this bin is the maximum. */
          if (localbuf[c + 1][i] > dmax) {
            /* store the new maximum index. */
            imax = i;
            dmax = localbuf[c + 1][imax];
          }
        }

        /* loop again through the local buffer frames. */
        for (i = 0; i < N; i++) {
          /* normalize by the maximum. */
          localbuf[c + 1][i] /= dmax;
        }
      }

      /* FIXME: yep! here's what we do here:
       *  1. rescale values in localbuf to [-1:1]
       *  2. change the timebase channel to [0:1]-scaled frequency values
       */

      /* attempt to write to the output pipe. */
      if (pipes_write (po, localbuf, N) != SP_OK) {
        /* print a failure message and break the loop. */
        warnf ("unable to write to output audio pipe");
        break;
      }

      /* loop through the remaining input buffer frames. */
      for (n = 0, i = k + 1; i < frames; i++) {
        /* store the timebase channel value in the local buffer. */
        localbuf[0][n] = rbuf[0][i];

        /* and through the channels of each frame. */
        for (c = 0; c < C; c++) {
          /* store the input channel value in the local buffer. */
          localbuf[c + 1][n] = rbuf[channels[c] - 1][i];
        }

        /* increment the frame count. */
        n++;
      }

      /* zero the cutoff index. */
      rollover = SP_FALSE;
      k = 0;
    }
    else {
      /* simply increment the size by the frame count and read again. */
      n += frames;
    }
  }

  /* free the input and output pipes buffer memories. */
  pipes_buf_free (pi, frames, rbuf);
  pipes_buf_free (po, frames, wbuf);

  /* free the local fft buffer memory. */
  pipes_buf_free (po, N, localbuf);

  /* attempt to close the input pipe. */
  if (pipes_close (pi) != SP_OK) {
    /* print a failure message and exit the application. */
    failf ("unable to close input pipe");
  }

  /* attempt to close the output pipe. */
  if (pipes_close (po) != SP_OK) {
    /* print a failure message and exit the application. */
    failf ("unable to close output pipe");
  }
}

